package com.inspur.edp.ai.flow.core.engine;

import com.inspur.edp.ai.flow.metadata.AbsNode;
import com.inspur.edp.ai.flow.metadata.FlowMetadata;
import com.inspur.edp.ai.flow.metadata.Parameter;
import com.inspur.edp.ai.flow.metadata.json.JsonType;
import com.inspur.edp.ai.flow.spi.GetOutputParameter;
import com.inspur.edp.ai.flow.spi.GetSubFlowNode;
import com.inspur.edp.ai.flow.spi.NodeCompileFactory;
import com.inspur.edp.ai.flow.spi.NodeCompiler;
import com.inspur.edp.common.expr.ast.ASTContext;
import com.inspur.edp.common.expr.ast.DefaultASTContext;
import com.inspur.edp.common.expr.evaluator.EvaluatorService;
import com.inspur.edp.common.expr.evaluator.ExpressEvaluator;
import com.inspur.edp.common.expr.visitor.field.HashMapFieldVisitor;
import com.inspur.edp.common.type.Field;
import com.inspur.edp.common.type.Type;
import com.inspur.edp.common.type.TypeRefer;
import com.inspur.edp.common.type.cache.TypeThreadCache;
import io.iec.edp.caf.commons.utils.StringUtils;

import java.util.*;

/**
 * @author lizhaorui
 * @date 2025/9/3
 * @description
 */

public class CompileEngine {

    public Map<String, Map<String, ExpressEvaluator>> buildAllNodeParams(FlowMetadata metadata) {
        Map<String, Map<String, ExpressEvaluator>> nodeInputParamMap = new HashMap<>();
        if (metadata.getNodes() == null) {
            return nodeInputParamMap;
        }

        List<AbsNode> allNodes = metadata.getAllNodes();

        // 1、遍历所有节点，初始化输出类型
        ASTContext astContext = new DefaultASTContext();
        for (AbsNode absNode : allNodes) {
            Map<String, TypeRefer> outputParamTypes = this.getOutputParamTypes(astContext, absNode);
            if (outputParamTypes.size() > 0) {
                Type nodeType = this.buildNodeOutputType(outputParamTypes);
                TypeThreadCache.add(nodeType);
                astContext.addVariable(absNode.getCode(), new TypeRefer(nodeType.getSource(), nodeType.getTypeId()));
            }
        }
        //2、遍历所有节点,编译节点
        for (AbsNode absNode : allNodes) {
            Map<String, ExpressEvaluator> nodeParamExpresses = this.buildNodeInputParamExpress(astContext, absNode);
            nodeInputParamMap.put(absNode.getId(), nodeParamExpresses);
        }
        return nodeInputParamMap;
    }

    private Map<String, TypeRefer> getOutputParamTypes(ASTContext astContext, AbsNode absNode) {
        Map<String, TypeRefer> outputTypeMap = new HashMap<>();
        NodeCompiler compile = NodeCompileFactory.get(absNode.getKind());
        List<Parameter> outputParameters = absNode.getOutputParams();
        if (compile != null && compile instanceof GetOutputParameter) {
            outputParameters = ((GetOutputParameter) compile).getOutputParams(astContext, absNode);
        }
        if (outputParameters != null) {
            for (Parameter outputParam : outputParameters) {
                TypeRefer outputType = null;
                if (outputParam.getSchema() == null) {
                    outputType = outputParam.getType();
                } else {
                    JsonType jsonType = outputParam.getSchema();
                    String typePrefix = String.format("%s_%s", absNode.getCode(), outputParam.getCode());
                    JsonTypeMgr.ConvertResult convertResult = JsonTypeMgr.convert(typePrefix, jsonType);
                    //添加到临时类型仓库
                    TypeThreadCache.addAll(convertResult.getTypes());
                    outputType = convertResult.getRootType();
                    outputParam.setType(outputType);
                }
                outputTypeMap.put(outputParam.getCode(), outputType);
            }
        }
        return outputTypeMap;
    }


    private Map<String, ExpressEvaluator> buildNodeInputParamExpress(ASTContext context, AbsNode node) {
        List<Parameter> inputParameters = node.getInputParams();
        NodeCompiler compiler = NodeCompileFactory.get(node.getKind());
        if (compiler == null) {
            Map<String, ExpressEvaluator> parameterMap = new HashMap<>();
            if (inputParameters != null) {
                for (Parameter parameter : inputParameters) {
                    ExpressEvaluator expressEvaluator = null;
                    if (!StringUtils.isEmpty(parameter.getValue())) {
                        expressEvaluator = EvaluatorService.build(context, parameter.getValue());
                    } else {
                        expressEvaluator = EvaluatorService.build(context, parameter.getValueExpr());
                    }
                    parameterMap.put(parameter.getCode(), expressEvaluator);
                }
            }
            return parameterMap;
        } else {
            return compiler.compile(context, node);
        }
    }

    private Type buildNodeOutputType(Map<String, TypeRefer> variableMap) {
        Type contextType = new Type();
        contextType.setTypeOptCode("java.util.Map");
        String uuid = UUID.randomUUID().toString().replace("-", "_");
        String typeId = "ContextType_" + uuid;
        contextType.setSource(TypeRefer.TEMP_SOURCE);
        contextType.setTypeId(typeId);
        contextType.setFields(new ArrayList<>());
        for (Map.Entry<String, TypeRefer> entry : variableMap.entrySet()) {
            Field contextField = new Field();
            contextField.setCode(entry.getKey());
            contextField.setType(entry.getValue());
            contextField.setAccessKind(HashMapFieldVisitor.KIND);
            contextType.getFields().add(contextField);
        }
        return contextType;
    }


}